package com.shujia.spark.opt

import org.apache.spark.SparkContext
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.SparkSession

object Demo2AggregateByKey {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local")
      .appName("cache")
      .config("spark.sql.shuffle.partitions", 1)
      .getOrCreate()

    val sc: SparkContext = spark.sparkContext
    val linesDS: RDD[String] = sc.textFile("data/students.txt")
    val kvRDD: RDD[(String, Int)] = linesDS.map((line: String) => {
      val split: Array[String] = line.split(",")
      (split(4), 1)
    })

    kvRDD
      .groupByKey()
      .map(kv => (kv._1, kv._2.size))
      .foreach(println)

    /**
     * 使用高性能算子
     * 使用reduceByKey代替groupByKey
     * reduceByKey会在map端做预聚合，可以减少shuffle过程中传说的数据量
     *
     * 大数据计算中shuffle时最费时间的，
     * 大部分的优化都是围绕shuffle做
     *
     */
    kvRDD
      .reduceByKey(_ + _)
      .foreach(println)


    /**
     * aggregateByKey: 需要指定三个参数
     * 1、初始值
     * 2、map端的聚合函数
     * 3、reduce端的聚合函数
     */
    kvRDD
      .aggregateByKey(0)(
        (u: Int, i: Int) => u + i, //map端的聚合函数
        (u1: Int, u2: Int) => u1 + u2 //reduce端聚合函数
      )
      .foreach(println)


    /**
     * 统计每个班级的平均年龄
     *
     */
    val clazzAndAgeRDD: RDD[(String, Int)] = linesDS.map((line: String) => {
      val split: Array[String] = line.split(",")
      (split(4), split(2).toInt)
    })
    clazzAndAgeRDD
      .groupByKey()
      .map {
        case (clazz: String, ages: Iterable[Int]) =>
          val avgAge: Double = ages.sum.toDouble / ages.size
          (clazz, avgAge)
      }
      .foreach(println)


    clazzAndAgeRDD
      //两个初始值，一个保存总人数，一个保存总的年龄
      .aggregateByKey((0, 0.0))(
        (kv: (Int, Double), age: Int) => (kv._1 + 1, kv._2 + age), //map端的聚合函数
        (kv1: (Int, Double), kv2: (Int, Double)) => (kv1._1 + kv2._1, kv1._2 + kv2._2) //reduce端聚合函数
      )
      .map {
        case (clazz: String, (num: Int, sumAge: Double)) =>
          (clazz, sumAge / num)
      }
      .foreach(println)


    while (true) {
    }
  }

}
